/*
Unpremultiply - Avisynth plugin to convert a premultipled RGB32 clip to straight matted

Description:
	This plugin takes no arguments other than the clip to be unpremultiplied. Example usage:
	videoclipname.Unpremultiply()
	This function will ONLY work on RGB32 input (and it really has no use on other image formats)

Notes:
	This plugin has only been tested on RGB32 AVI files exported from Eyeon Fusion compositing software.
	Using it on PNG files from the same software has typically resulted in garbage.
	This plugin is based almost entirely on the FilterSDK examples from avisynth.org

* Copyright (c) 2008, Josh Sutinen <josh@damagedgoodz.net>
* All rights reserved.
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*     * Redistributions of source code must retain the above copyright
*       notice, this list of conditions and the following disclaimer.
*     * Redistributions in binary form must reproduce the above copyright
*       notice, this list of conditions and the following disclaimer in the
*       documentation and/or other materials provided with the distribution.
*     * Neither the name of the organization nor the
*       names of its contributors may be used to endorse or promote products
*       derived from this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDER ''AS IS'' AND ANY
* EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
* WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER BE LIABLE FOR ANY
* DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
* (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
* LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
* ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
* SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

#include "windows.h"
#include "avisynth.h"
#include "stdint.h"

class Unpremultiply : public GenericVideoFilter 
{
public:
	Unpremultiply(PClip _child, IScriptEnvironment* env);
	PVideoFrame __stdcall GetFrame(int n, IScriptEnvironment* env);
};

Unpremultiply::Unpremultiply(PClip _child, IScriptEnvironment* env) : GenericVideoFilter(_child) 
{
	// Constructor goes in here
	// Following variables are defined:
	// PClip child; // Contains source clip
	// VideoInfo vi; // Contains videoinfo from the source clip

	// This filter is designed strictly for removing black premultiplied matting from overlay clips
	// Therefore we require RGB32 (with alpha channel)
	if(!vi.IsRGB32())
		env->ThrowError("Unpremultiply: input to filter must be in RGB32 colorspace");
}

PVideoFrame __stdcall Unpremultiply::GetFrame(int n, IScriptEnvironment* env)
{
	// Get frame n from the source clip
	PVideoFrame src = child->GetFrame(n, env);
	// Construct a new frame based on the info of the current frame
	PVideoFrame dst = env->NewVideoFrame(vi);

	// Get a Read pointer from the current source frame
	// For RGB sources this will point to the lower left pixel
	const unsigned char* srcp = src->GetReadPtr();

	// Get a Write pointer to the destination frame
	unsigned char* dstp = dst->GetWritePtr();

	// Get pitch (line length) of the destination image
	// Pitch will always be equal to or greater than the width in pixels
	const int dst_pitch = dst->GetPitch();

	// Get rowsize (number of used bytes in a line)
	const int dst_width = dst->GetRowSize();

	// Get height of destination image
	const int dst_height = dst->GetHeight();

	// Get pitch, width, and height of the source frame
	const int src_pitch = src->GetPitch();
	const int src_width = src->GetRowSize();
	const int src_height = src->GetHeight();

	int w,h;

	for(h=0; h < src_height; h++)
	{
		for(w=0; w< src_width; w+=4)
		{
			uint8_t srcbpix = *(srcp + w);
			uint8_t srcgpix = *(srcp + w + 1);
			uint8_t srcrpix = *(srcp + w + 2);
			uint8_t srcapix = *(srcp + w + 3);
			uint8_t* dstbpix = (dstp + w);
			uint8_t* dstgpix = (dstp + w + 1);
			uint8_t* dstrpix = (dstp + w + 2);
			if(srcapix == 0)
			{
				*dstbpix = srcbpix;
				*dstgpix = srcgpix;
				*dstrpix = srcrpix;
			} else
			{
				// Formula is: Cu = Cm / a
				// Adding alpha/2 rounds instead of truncates the result of integer division
				// This will cause horrible things to happen to your image if a color value
				// is ever greater than the alpha value.
				// This can happen if your source program used some kind of specular highlighting
				*dstbpix = (srcbpix * 255 + srcapix / 2) / srcapix;
				*dstgpix = (srcgpix * 255 + srcapix / 2) / srcapix;
				*dstrpix = (srcrpix * 255 + srcapix / 2) / srcapix;
			}
			// Alpha is a direct copy
			*(dstp + w + 3) = *(srcp + w + 3);
		}

		// Add the pitch of one line (in bytes) to the source pointer and destination pointer
		srcp = srcp + src_pitch;
		dstp = dstp + dst_pitch;
	}

	// Done processing the image, return the result image
	return dst;
}

// This is the function that created the filter, when the fitler has been called
AVSValue __cdecl Create_Unpremultiply(AVSValue args, void* user_data, IScriptEnvironment* env)
{
	return new Unpremultiply(args[0].AsClip(),env);
}

// The following function actually registers the filter in Avisynth
// It is called automagically when the plugin is loaded to see which functions this filter contains
extern "C" __declspec(dllexport) const char* __stdcall AvisynthPluginInit2(IScriptEnvironment* eng)
{
	eng->AddFunction("Unpremultiply", "c", Create_Unpremultiply,0);
	return "Unpremultiply - Alpha matte remover";
}